package ssh

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"log"
	"regexp"
	"strconv"
	"time"

	"gitee.com/leminewx/gokit"
	"gitee.com/leminewx/gokit/util/pool"
	"golang.org/x/crypto/ssh"
)

const (
	DEFAULT_PORT    = 22
	DEFAULT_TIMEOUT = 20
)

// Client 定义 SSH 协议客户端的数据结构
type Client struct {
	debug  bool
	port   int
	host   string
	addr   string
	config *ssh.ClientConfig
	client *ssh.Client
}

// NewClient 新建一个 SSH 客户端
func NewClient(host, user, pwd string, port ...int) (*Client, error) {
	client := &Client{
		host: host,
		port: DEFAULT_PORT,
		addr: fmt.Sprintf("%s:%d", host, DEFAULT_PORT),
		config: &ssh.ClientConfig{
			User:            user,
			Auth:            []ssh.AuthMethod{ssh.Password(pwd)},
			Timeout:         DEFAULT_TIMEOUT * time.Second,
			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		},
	}

	if len(port) > 0 {
		client.port = port[0]
		client.addr = fmt.Sprintf("%s:%d", client.host, client.port)
	}

	var err error
	client.client, err = ssh.Dial("tcp", client.addr, client.config)
	return client, err
}

// WithDebug 开启调试模式
func (own *Client) WithDebug() *Client {
	own.debug = true
	return own
}

// Execute 执行命令，并返回结果
func (own *Client) Execute(cmd string) (string, error) {
	// 创建会话
	session, err := own.client.NewSession()
	if err != nil {
		return "", err
	}

	// 执行命令
	resp, err := session.CombinedOutput(cmd)
	if own.debug {
		log.Printf("[%s] [SSH] Execute SSH command: %s\n%s\n", own.addr, cmd, resp)
	}

	if exitErr, ok := IsExitError(err); ok {
		switch exitErr.ExitStatus() {
		case 1, 141, 127:
			return string(resp), nil
		}
	}

	return TrimASNIBytes(resp), err
}

// ExecuteWithTimeout 执行命令，并指定超时时间
func (own *Client) ExecuteWithTimeout(cmd string, timeout time.Duration) (string, error) {
	// 创建会话
	session, err := own.client.NewSession()
	if err != nil {
		return "", err
	}

	// 执行命令
	var buf bytes.Buffer
	session.Stdout, session.Stderr = &buf, &buf
	if err = session.Start(cmd); err != nil {
		session.Close() // 关闭会话
		return "", err
	}

	// 异步等待正常结束
	done := make(chan error, 1)
	go func() {
		done <- session.Wait()
	}()

	// 同步等待会话结束
	select {
	case err = <-done: // 执行结束
	case <-time.After(timeout): // 执行超时
		// 关闭会话，并等待 100ms
		session.Close()
		<-time.After(100 * time.Millisecond)

		// 等待会话结束
		select {
		case err = <-done:
		default: // 强制结束
		}
	}

	// 获取结果
	resp := buf.Bytes()
	if own.debug {
		log.Printf("[%s] [SSH] Execute SSH command: %s\n%s", own.addr, cmd, resp)
	}

	// 检查是否为空
	if exitErr, ok := IsExitError(err); ok {
		switch exitErr.ExitStatus() {
		case 1, 141, 127:
			return string(resp), nil
		}
	}

	return TrimASNIBytes(resp), err
}

// ExecuteWithTimeout 执行命令，并通过上下文控制
func (own *Client) ExecuteWithContext(cmd string, ctx context.Context) (string, error) {
	if ctx == nil {
		return "", fmt.Errorf("gokit: ssh: nil context")
	}

	// 创建会话
	session, err := own.client.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()

	// 执行命令
	var buf bytes.Buffer
	session.Stdout, session.Stderr = &buf, &buf
	if err = session.Start(cmd); err != nil {
		return "", err
	}

	// 异步等待正常结束
	done := make(chan error, 1)
	go func() {
		done <- session.Wait()
	}()

	// 同步等待会话结束
	select {
	case err = <-done: // 执行结束
	case <-ctx.Done(): // 执行超时
		// 关闭会话，并等待 100ms
		session.Close()
		<-time.After(100 * time.Millisecond)

		// 等待会话结束
		select {
		case err = <-done:
		default: // 强制结束
		}
	}

	// 等待结束
	resp := buf.Bytes()
	if own.debug {
		log.Printf("[%s] [SSH] Execute SSH command: %s\n%s", own.addr, cmd, resp)
	}

	if exitErr, ok := IsExitError(err); ok {
		switch exitErr.ExitStatus() {
		case 1, 141, 127:
			return string(resp), nil
		}
	}

	return TrimASNIBytes(resp), err
}

// PTY 定义伪终端参数的数据结构
type PTY struct {
	// Term 终端类型
	Term string
	// Height, Width 终端的高度和宽度
	Height, Width int
	// Modes 终端样式
	Modes ssh.TerminalModes
	// Timeout 执行超时时间
	Timeout time.Duration
}

var DefaultPTY = PTY{
	Term:   "xterm",
	Height: 150,
	Width:  50,
	Modes: ssh.TerminalModes{
		ssh.ECHO:          0,     // 禁用回显（用于密码输入）
		ssh.TTY_OP_ISPEED: 14400, // 输入速度
		ssh.TTY_OP_OSPEED: 14400, // 输出速度
	},
	Timeout: 30 * time.Second,
}

// ExecuteWithPTY 执行交互式命令，并通过伪终端实现
func (own *Client) ExecuteWithPTY(cmds []string, pty ...*PTY) (string, error) {
	var _pty *PTY
	if len(pty) > 0 {
		_pty = pty[0]
	} else {
		p := DefaultPTY
		_pty = &p
	}

	// 创建会话
	session, err := own.client.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()

	// 发起伪终端请求
	if err = session.RequestPty(_pty.Term, _pty.Height, _pty.Width, _pty.Modes); err != nil {
		return "", err
	}

	// 创建管道
	stdinPipe, _ := session.StdinPipe()
	stdoutPipe, _ := session.StdoutPipe()
	stderrPipe, _ := session.StderrPipe()

	// 启动shell
	if err = session.Shell(); err != nil {
		return "", err
	}

	// 异步获取输出
	outBuffer := pool.DefaultBytesBufferPool.Alloc()
	go func() {
		scanner := bufio.NewScanner(stdoutPipe)
		for scanner.Scan() {
			outBuffer.WriteString(TrimASNIBytes([]byte(scanner.Text())))
			outBuffer.WriteByte('\n')
		}
	}()

	// 异步处理错误输出
	errBuffer := pool.DefaultBytesBufferPool.Alloc()
	go func() {
		scanner := bufio.NewScanner(stderrPipe)
		for scanner.Scan() {
			errBuffer.WriteString(TrimASNIBytes([]byte(scanner.Text())))
			errBuffer.WriteByte('\n')
		}
	}()
	defer func() {
		pool.DefaultBytesBufferPool.Dealloc(outBuffer)
		pool.DefaultBytesBufferPool.Dealloc(errBuffer)
	}()

	// 异步执行交互命令
	go func() {
		if cmds[len(cmds)-1] != "exit" {
			cmds = append(cmds, "exit")
		}

		time.Sleep(500 * time.Millisecond) // 等待shell初始化
		for _, cmd := range cmds {
			stdinPipe.Write([]byte(cmd + "\n"))
			time.Sleep(100 * time.Millisecond)
		}
	}()

	// 等待会话结束
	if err = session.Wait(); err != nil {
		if exitErr, ok := IsExitError(err); ok {
			switch exitErr.ExitStatus() {
			case 1, 141, 127:
				return TrimASNIBytes(outBuffer.Bytes()), nil
			}
		}
	}

	resp := TrimASNIBytes(outBuffer.Bytes())
	if own.debug {
		log.Printf("[%s] [SSH] Execute SSH command: %v\n%s", own.addr, cmds, resp)
	}

	return resp, err
}

// Close 关闭 SSH 客户端
func (own *Client) Close() error {
	return own.client.Close()
}

// 编译正则表达式，匹配所有 ANSI 转义码
var ansiRegexp = regexp.MustCompile(`\x1B\[[0-?]*[ -/]*[@-~]`)

// TrimASNIBytes 剪除所有 ANSI 转义符
func TrimASNIBytes(data []byte) string {
	return string(ansiRegexp.ReplaceAll(data, []byte("")))
}

// IsExitError 检查错误是否是 Exit 错误
func IsExitError(err error) (*ssh.ExitError, bool) {
	e, ok := err.(*ssh.ExitError)
	return e, ok
}

// FormatStringSliceToIntSlice 将数字型字符串列表转换为整型列表
func FormatStringSliceToIntSlice(items []string) []int {
	res := make([]int, 0, len(items))
	for _, pid := range items {
		p, _ := strconv.ParseInt(pid, 10, 64)
		res = append(res, int(p))
	}

	return res
}

// JoinNumberSliceToString 拼接整型列表为字符串
func JoinNumberSliceToString[T gokit.Number](items []T, sep string) string {
	var res string
	first := true
	for _, item := range items {
		if !first {
			res += sep
		}

		res += strconv.Itoa(int(item))
		first = false
	}

	return res
}
